﻿using System;
using System.Collections.Generic;
using System.Text;
using ESRI.ArcGIS.DataSourcesRaster;
using ESRI.ArcGIS.Geodatabase;
using ESRI.ArcGIS.DataSourcesGDB;
using System.Diagnostics;
using ESRI.ArcGIS.SpatialAnalyst;


namespace CanSWAT.ArcSWAT
{

    enum SlopeType
    {
        Slope,
        Curvature
    }


    class RasterDatabase : BaseWithProcess
    {
        private static string RASTER_DATABASE_NAME = "RasterStore.mdb"; //raster database name
        private static string SOURCE_DEM_NAME = "sourcedem";            //source dem, used to calculate slope and curvature when the slope is bigger than clip dem
        private static string CLIP_DEM_NAME = "ClipDEM";         //clip dem, used to calculate slope and curvature when the slope is smaller than source dem
        private static string SLOPE_ORIGINAL_NAME = "slope";            //slope
        private static string SLOPE_NAME = "slope_canswat";             //slope along wind direction
        private static string CURVATURE_NAME = "curvature_canswat";     //curvature along wind direction
        private static string WATERSHED_NAME = "Watershed";
        private static int[] DX = { 1, 1, 0, -1, -1, -1, 0, 1 };
        private static int[] DY = { 0, 1, 1, 1, 0, -1, -1, -1 };
        private static byte DEFAULT_WATERSHED_NO_DATA = 0;

        private string _rasterDatabasePath = "";                        //raster database path

        private IRasterWorkspaceEx _workspace = null;                   //raster workspace
        private IRasterDataset _originalSlopeDataset = null;            //raster dataset, similar to featureclass dataset
        private IRasterDataset _slopeDataset = null;
        private IRasterDataset _curvatureDataset = null;
        private IRasterDataset _sourceDEMDataset = null;

        private IRaster _sourceDEM = null;
        private IRaster _slope = null;
        private IRaster _curvature = null;      

        private System.Array _sourceDEM_BlockData = null;
        private System.Array _slope_BlockData = null;
        private System.Array _watershed_BlockData = null;
        private int _dem_width = -1;
        private int _dem_height = -1;
        private double _cellSize = -1.0;

        private Project _parentProject = null;

        public RasterDatabase(string projectFolder, Project prj)
        {
            string[] mdbFiles = System.IO.Directory.GetFiles(projectFolder, RASTER_DATABASE_NAME);
            if (mdbFiles.Length > 0)
                _rasterDatabasePath = mdbFiles[0];

            _parentProject = prj;
        }

        public bool IsValid
        {
            get { return _rasterDatabasePath.Length > 0 && SourceDEMDataset != null && OriginalSlopeDataset != null; }
        }

        ~RasterDatabase()
        {
            if(_sourceDEM != null)
                System.Runtime.InteropServices.Marshal.ReleaseComObject(_sourceDEM);
            if (_slope != null)
                System.Runtime.InteropServices.Marshal.ReleaseComObject(_slope);
            if (_curvature != null)
                System.Runtime.InteropServices.Marshal.ReleaseComObject(_curvature);
        }

        private IRasterWorkspaceEx Workspace
        {
            get
            {
                if (_rasterDatabasePath.Length == 0) return null;

                if (_workspace == null)
                {
                    IWorkspaceFactory wf = new AccessWorkspaceFactoryClass();
                    _workspace = wf.OpenFromFile(_rasterDatabasePath, 0) as IRasterWorkspaceEx;
                }
                return _workspace;
            }
        }

        private IRasterDataset findRasterDataset(string name)
        {
            if (Workspace == null) return null;

            IWorkspace2 rw2 = Workspace as IWorkspace2;
            if (!rw2.get_NameExists(esriDatasetType.esriDTRasterDataset, name))
            {
                //copy and create slope and curvature raster dataset if they don't exist
                if (name.ToLower().Equals(SLOPE_NAME) || name.ToLower().Equals(CURVATURE_NAME))
                {
                    if (OriginalSlopeDataset != null && OriginalSlopeDataset.CanCopy())
                    {
                        System.Diagnostics.Debug.WriteLine("Creating raster dataset " + name + " ... " + DateTime.Now.ToString());
                        OriginalSlopeDataset.Copy(name, Workspace as IWorkspace);
                        System.Diagnostics.Debug.WriteLine("Creating raster dataset " + name + " finished. " + DateTime.Now.ToString());
                    }
                    else
                        return null;
                }
                else return null;
            }
        
            return this.Workspace.OpenRasterDataset(name);
        }

        private IRaster getRasterFromDataset(IRasterDataset ds)
        {
            if (ds == null) return null;

            Debug.WriteLine("Creating default raster for " +(ds as IDataset).Name + " ... " + DateTime.Now.ToString());
            IRaster r = ds.CreateDefaultRaster();
            Debug.WriteLine("Creating default raster finished. " + DateTime.Now.ToString());

            return r;
        }

        private IRasterDataset OriginalSlopeDataset
        {
            get
            {
                if (_originalSlopeDataset == null)
                    _originalSlopeDataset = findRasterDataset(SLOPE_ORIGINAL_NAME);

                return _originalSlopeDataset;
            }
        }

        private IRasterDataset SlopeDataset
        {
            get
            {
                if (_slopeDataset == null)
                    _slopeDataset = findRasterDataset(SLOPE_NAME);

                return _slopeDataset;
            }
        }

        private IRasterDataset CurvatureDataset
        {
            get
            {
                if (_curvatureDataset == null)
                    _curvatureDataset = findRasterDataset(CURVATURE_NAME);

                return _curvatureDataset;
            }
        }

        private IRasterDataset SourceDEMDataset
        {
            get
            {
                if (_sourceDEMDataset == null)
                {
                    //For La Sale watershed, the slope is in same size of sourceDEM.
                    //For Yongbo's sheep river watershed, the slope is in same size of clipDEM.
                    //Slope ans Watershed is always in same size.
                    //compare the size of slope and clip dem and source dem to decide which dem will be used.
                    IRasterProps props = Slope as IRasterProps;
                    
                    IRasterDataset sourceDemDataset = findRasterDataset(SOURCE_DEM_NAME);
                    IRaster sourceDem = getRasterFromDataset(sourceDemDataset);
                    IRasterProps sourceDemProps = sourceDem as IRasterProps;
                    if (sourceDemProps.Height == props.Height &&
                        sourceDemProps.Width == props.Width)
                    {
                        _sourceDEMDataset = sourceDemDataset;
                        _sourceDEM = sourceDem;
                        CanSWAT.debugInformation("SourceDEM is used.");
                    }
                    else
                    {
                        //release source dem first
                        System.Runtime.InteropServices.Marshal.ReleaseComObject(sourceDem);

                        IRasterDataset clipDemDataset = findRasterDataset(CLIP_DEM_NAME);
                        IRaster clipDem = getRasterFromDataset(clipDemDataset);
                        IRasterProps clipDemProps = clipDem as IRasterProps;

                        if (clipDemProps.Height == props.Height &&
                            clipDemProps.Width == props.Width)
                        {
                            _sourceDEMDataset = clipDemDataset;
                            _sourceDEM = clipDem;
                            CanSWAT.debugInformation("ClipDEM is used.");
                        }
                        else
                        {
                            System.Runtime.InteropServices.Marshal.ReleaseComObject(clipDem);
                            throw new System.Exception("The size of slope/watershed is different with SourceDEM and ClipDEM. Please check you project.");
                        }
                    }

                }

                return _sourceDEMDataset;
            }
        }

        private IRaster SourceDEM
        {
            get 
            {
                if (_sourceDEM == null)
                    _sourceDEM = getRasterFromDataset(SourceDEMDataset);

                return _sourceDEM;
            }
        }

        private IRaster Slope
        {
            get
            {
                if (_slope == null)
                    _slope = getRasterFromDataset(SlopeDataset);

                return _slope;
            }
        }

        private IRaster Curvature
        {
            get
            {
                if (_curvature == null)
                    _curvature = getRasterFromDataset(CurvatureDataset);

                return _curvature;
            }
        }

        private int DEMWidth
        {
            get
            {
                if (_dem_width < 0)
                {
                    IRasterProps props = SourceDEM as IRasterProps;
                    _dem_width = props.Width;
                }
                return _dem_width;
            }
        }

        private int DEMHeight
        {
            get
            {
                if (_dem_height < 0)
                {
                    IRasterProps props = SourceDEM as IRasterProps;
                    _dem_height = props.Height;
                }
                return _dem_height;
            }
        }

        /// <summary>
        /// All raster has same cell size
        /// </summary>
        private double CellSize
        {
            get
            {
                if (_cellSize < 0)
                    _cellSize = (SourceDEM as IRasterProps).MeanCellSize().X;
                return _cellSize;
            }
        }

        /// <summary>
        /// read all the raster data in first band from given raster. 
        /// </summary>
        /// <param name="r">raster data</param>
        /// <returns>raster data in 2-d array format</returns>
        /// <remarks>The given raster must has the same size with source DEM, such as new slope and curvature</remarks>
        private System.Array readRasterData(IRaster r)
        {
            return readRasterData(r, null);
        }

        /// <summary>
        /// read all the raster data in first band from given raster. 
        /// </summary>
        /// <param name="r">raster data</param>
        /// <param name="isSetNoData">if set a default value to no data cells</param>
        /// <returns>raster data in 2-d array format</returns>
        /// <remarks>The given raster must has the same size with source DEM, such as new slope and curvature</remarks>
        private System.Array readRasterData(IRaster r, object nodata)
        {
            IPnt blocksize = new PntClass();
            blocksize.SetCoords(DEMWidth, DEMHeight);

            IPnt upperLeft = new PntClass();
            upperLeft.SetCoords(0, 0);

            //create just one block
            //all cell values are 0 at this point, IRaster.Read() needs to be called to populate the block with the real data.
            IPixelBlock3 block = r.CreatePixelBlock(blocksize) as IPixelBlock3;

            r.Read(upperLeft, block as IPixelBlock);

            System.Array data = (System.Array)block.get_PixelData(0);

            if (nodata == null) return data;

            //set a default no data to all the no data cells
            //this just change the memory, will not affect the data in disk
            //used to indentify no data cell easily
            //only used for waterhsed data in elevation band calculation
            for (int i = 0; i < DEMWidth; i++)
            {
                for (int j = 0; j < DEMHeight; j++)
                {
                    //if current cell is not no data, do nothing
                    if (Convert.ToInt32(block.GetNoDataMaskVal(0, i, j).ToString()) == 1)
                        continue;

                    data.SetValue(nodata,i,j);
                }
            }
            return data;
        }

        /// <summary>
        /// retrieve elevation of given cell from Source DEM for slope calculation.
        /// </summary>
        /// <param name="col">The column of the raster</param>
        /// <param name="row">The row of the raster</param>
        /// <returns></returns>
        private double getElevation(int col, int row)
        {
            if (col < 0 || col >= DEMWidth) return -1.0;
            if (row < 0 || row >= DEMHeight) return -1.0;

            if (_sourceDEM_BlockData == null)
                _sourceDEM_BlockData = readRasterData(SourceDEM);

            return Convert.ToDouble(_sourceDEM_BlockData.GetValue(col, row).ToString()); 
        }

        /// <summary>
        /// retrieve slope of given cell fro slope raster for curvature calculation
        /// </summary>
        /// <param name="col">The column of the raster</param>
        /// <param name="row">The row of the raster</param>
        /// <returns></returns>
        private double getSlope(int col, int row)
        {
            if (col < 0 || col >= DEMWidth) return -1.0;
            if (row < 0 || row >= DEMHeight) return -1.0;

            if (_slope_BlockData == null)
                _slope_BlockData = readRasterData(Slope);

            return Convert.ToDouble(_slope_BlockData.GetValue(col, row).ToString());
        }

        /// <summary>
        /// retrieve elevation (for slope calculation) or slope (for curvature calculation)
        /// </summary>
        /// <param name="col">The column of the raster</param>
        /// <param name="row">The row of the raster</param>
        /// <param name="type">slope or curvature</param>
        /// <returns></returns>
        private double getElevation(int col, int row, SlopeType type)
        {
            if (type == SlopeType.Slope) return getElevation(col, row);
            else return getSlope(col, row);
        }

        /// <summary>
        /// update all data corresponding to wind direction, including raster, FullHRU and hrus.
        /// </summary>
        /// <param name="windDirection">The wind direction. Default is 3 northwest. 
        /// <para>0 - East</para>
        /// <para>1 - Northeast</para>
        /// <para>2 - North</para>
        /// <para>3 - Northwest</para>
        /// <para>4 - West</para>
        /// <para>5 - Southwest</para>
        /// <para>6 - South</para>
        /// <para>7 - Southeast</para>
        /// </param>
        ///<remarks>runs very slow when working with background worker.</remarks>
        public void updateSlopeAndCurvature(int windDirection)
        {
            if(!IsValid)
                throw new System.Exception("Current raster database in not valid.");

            if (_parentProject == null || _parentProject.ProjectDatabase == null)
                throw new System.Exception("The ProjectDatabase object is null.");

            //if (prj.WindDirection == windDirection)
            //{
            //    System.Windows.Forms.MessageBox.Show("The wind direction is same with previous setting. No operation will proceed.", CanSWAT.CAN_SWAT_TITLE);
            //    return;
            //}

            //update raster dataset      
            setProgress(0, "Calculating slope and curvatrue along selected wind direction...");
            updateSlopeAndCurvatureRaster(windDirection);

            //put slope and curvature into FullHRU
            setProgress(40, "Calculating slope along selected wind direction for each HRU...");
            _parentProject.ProjectDatabase.calSlope(this, SlopeType.Slope, true);

            setProgress(60, "Calculating curvature along selected wind direction for each HRU...");
            _parentProject.ProjectDatabase.calSlope(this, SlopeType.Curvature, true);

            //put slope and curvature into hrus
            setProgress(80, "Saving slope and curvature along selected wind direction into table...");
            _parentProject.ProjectDatabase.calHRUParameter(true);

            //update wind direction
            _parentProject.ProjectDatabase.WindDirection = windDirection;
            setProgress(100, "DONE!");
        }        

        /// <summary>
        /// calculate/update slope and curvature along wind direction.
        /// create raster dataset if necessary. 
        /// </summary>
        /// <param name="windDirection">The wind direction. Default is 3 northwest. 
        /// <para>0 - East</para>
        /// <para>1 - Northeast</para>
        /// <para>2 - North</para>
        /// <para>3 - Northwest</para>
        /// <para>4 - West</para>
        /// <para>5 - Southwest</para>
        /// <para>6 - South</para>
        /// <para>7 - Southeast</para>
        /// </param>
        /// <remarks>slope must be updated before curvature</remarks>
        private void updateSlopeAndCurvatureRaster(int windDirection)
        {
            updateSlope(windDirection, SlopeType.Slope);
            updateSlope(windDirection, SlopeType.Curvature);
        }

        /// <summary>
        /// calculate/update slope and curvature along wind direction.
        /// create raster dataset if necessary. 
        /// </summary>
        /// <param name="windDirection">The wind direction. Default is 3 northwest. 
        /// <param name="type">slope or curvature</param>
        /// <para>0 - East</para>
        /// <para>1 - Northeast</para>
        /// <para>2 - North</para>
        /// <para>3 - Northwest</para>
        /// <para>4 - West</para>
        /// <para>5 - Southwest</para>
        /// <para>6 - South</para>
        /// <para>7 - Southeast</para>
        /// </param>
        /// <remarks>This is the core codes to calculate slope and curvature</remarks>
        private void updateSlope(int windDirection,SlopeType type)
        {
            if (SourceDEM == null)
                throw new System.Exception("Can't find SourceDEM.");
            
            //determin which is the destiniation raster
            IRaster destinationRaster = null;
            if (type == SlopeType.Slope) destinationRaster = Slope;
            else destinationRaster = Curvature;

            if (destinationRaster == null)
                throw new System.Exception("Could't create " + type.ToString());

            if (type == SlopeType.Slope)
            {
                Debug.WriteLine("Calculating slope along wind direction..." + DateTime.Now.ToString());
                setProgress(5, "Calculating slope along wind direction...");
            }
            else
            {
                Debug.WriteLine("Calculating Curvature along wind direction..." + DateTime.Now.ToString());
                setProgress(20, "Calculating slope along wind direction...");
            }

            //block size
            IPnt blocksize = new PntClass();
            blocksize.SetCoords(DEMWidth, DEMHeight);

            //create just one block
            //before IPixelBlock3.Read() is called, all value in the block is 0.
            IPixelBlock3 block = destinationRaster.CreatePixelBlock(blocksize) as IPixelBlock3;
            
            //Define the location that the upper left corner of the pixel block is to write.
            IPnt upperLeft = new PntClass();
            upperLeft.SetCoords(0, 0);

            //read data to block
            destinationRaster.Read(upperLeft, block as IPixelBlock);

            //get a array with the same size of block            
            System.Array pixels = (System.Array)block.get_PixelData(0);
            
            int opwinddir; //opposite wind direction
            if (windDirection > 3) 
                opwinddir = windDirection - 4;
            else 
                opwinddir = windDirection + 4;

            //temp variables
            int x0, y0, x1, y1, x2, y2;
            double elev0, elev1, elev2, finalslope, slope1, slope2;


            for (int i = 0; i < DEMWidth; i++)
            {
                for (int j = 0; j < DEMHeight; j++)
                {
                    //if current cell is no data, do nothing
                    if(Convert.ToInt32(block.GetNoDataMaskVal(0,i,j).ToString()) == 0)
                        continue;

                    x1 = i;
                    y1 = j;
                    x0 = x1 + DX[opwinddir];
                    y0 = y1 + DY[opwinddir];
                    x2 = x1 + DX[windDirection];
                    y2 = y1 + DY[windDirection];

                    elev0 = getElevation(x0, y0, type);
                    elev1 = getElevation(x1, y1, type);
                    elev2 = getElevation(x2, y2, type);

                    if (elev0 > 0 && elev1 > 0)
                        slope1 = (elev1 - elev0) / CellSize;
                    else
                        slope1 = 0.0;

                    if (elev1 > 0 && elev2 > 0)
                        slope2 = (elev2 - elev1) / CellSize;
                    else
                        slope2 = 0.0;

                    if (slope1 == 0) finalslope = slope2;
                    else if (slope2 == 0) finalslope = slope1;
                    else finalslope = slope1 / 2.0 + slope2 / 2.0;

                    //set the slope
                    pixels.SetValue(Convert.ToSingle(finalslope * 100.0), i, j);
                }
            }

            //set the new data to raster dataset
            block.set_PixelData(0, pixels);

            //Write the pixel block.
            IRasterEdit rasterEdit = (IRasterEdit)destinationRaster;
            rasterEdit.Write(upperLeft, (IPixelBlock)block);

            //must release here, otherwise error will happen when the dataset is read next time
            System.Runtime.InteropServices.Marshal.ReleaseComObject(rasterEdit);
            if (type == SlopeType.Slope) _slope = null;
            else _curvature = null;

            //set current slope block data to null to be updated when used next time
            if (type == SlopeType.Slope && _slope_BlockData != null)
            {
                System.Runtime.InteropServices.Marshal.ReleaseComObject(_slope_BlockData);
                _slope_BlockData = null;
            }

            if (type == SlopeType.Slope)
                Debug.WriteLine("Calculating slope along wind direction finished. " + DateTime.Now.ToString());
            else
                Debug.WriteLine("Calculating Curvature along wind direction finished." + DateTime.Now.ToString());
         }

        /// <summary>
        /// calculate slope/curvature aling with wind direction using spatial analysis extension.
        /// The spatial analysis exetension should be enabled in ArcMap to use this function. It may needs
        /// around 20 seconds.
        /// </summary>
        /// <param name="hru">FullHRU feature class which is used to define the zone</param>
        /// <param name="type">slope or curvature</param>
        /// <returns>It's a table populated with all the statistical values of each zone defined in FullHRU. 
        /// Only mean is used. So performance coould be improved if only mean is calculated.</returns>
        ///<remarks>This zonal statistic operation can't be executed in multi-thread program.</remarks>
        public ITable calSlope(IFeatureClass hru, SlopeType type)
        {
            if (!IsValid)
                throw new System.Exception("Current raster database in not valid.");

            IGeoDataset zone = hru as IGeoDataset;
            IGeoDataset value = null;
            string slopeName = "slope";

            if (type == SlopeType.Slope)
                value = SlopeDataset as IGeoDataset;
            else
            {
                value = CurvatureDataset as IGeoDataset;
                slopeName = "curvature";
            }

            Debug.WriteLine("Caculating mean "+slopeName+" for each HRU ..." + DateTime.Now.ToString());
            
            IZonalOp op = new RasterZonalOpClass();        
            ITable t = op.ZonalStatisticsAsTable(zone,value, true); //this function must be executed in main thread.
            
            Debug.WriteLine("Caculating mean " + slopeName + " for each HRU finished. " + DateTime.Now.ToString());

            return t;
        }

        /// <summary>
        /// get subbasin id for given cell
        /// </summary>
        /// <param name="col"></param>
        /// <param name="row"></param>
        /// <returns></returns>
        private int getSubbasinID(int col, int row)
        {
            //read source DEM data if necessary
            if (_sourceDEM_BlockData == null)
                _sourceDEM_BlockData = readRasterData(SourceDEM);

            //read subbain raster data
            //a default no data value is set to all no data cells to indentify no data cell
            if (_watershed_BlockData == null)
            {
                IRaster watershedRaster = getRasterFromDataset(findRasterDataset(WATERSHED_NAME));
                _watershed_BlockData = this.readRasterData(watershedRaster,DEFAULT_WATERSHED_NO_DATA);
                System.Runtime.InteropServices.Marshal.ReleaseComObject(watershedRaster);
            }

            return Convert.ToInt32(_watershed_BlockData.GetValue(col, row).ToString());               
        }


        public void getSubbasinElevation(SortedDictionary<int, Subbasin> subbasins)
        {
            CanSWAT.debugInformation("Reading elevation of each subbasin...");
            for (int i = 0; i < DEMWidth; i++)
            {
                for (int j = 0; j < DEMHeight; j++)
                {
                    int subid = getSubbasinID(i, j);

                    //if current cell is no data, do nothing
                    if (subid == DEFAULT_WATERSHED_NO_DATA) 
                        continue;

                    if (!subbasins.ContainsKey(subid))
                    {
                        Debug.WriteLine("Warning: Can't find subbasin " + subid.ToString());
                        continue;
                    }

                    Subbasin sub = subbasins[subid];
                    sub.AddElevation(Convert.ToInt32(getElevation(i, j)));
                }
            }
            CanSWAT.debugInformation("Reading elevation of each subbasin finished.");

        }
    }

    
}
